import os
import random
import shutil

def split_dataset(srcDir, trainDir, valDir, split_ratio=0.9):
    """
    将数据集划分为训练集和验证集，并保存到相应的文件夹中。

    Parameters:
    - srcDir: 原始数据集文件夹路径，包含图像和标签文件。
    - trainDir: 训练集文件夹路径，包含 'images' 和 'labels' 子文件夹。
    - valDir: 验证集文件夹路径，包含 'images' 和 'labels' 子文件夹。
    - split_ratio: 数据集划分比例，默认为 0.9，表示将 90% 的数据用于训练集，10% 用于验证集。
    """
    os.makedirs(os.path.join(trainDir, 'images'), exist_ok=True)
    os.makedirs(os.path.join(trainDir, 'labels'), exist_ok=True)
    os.makedirs(os.path.join(valDir, 'images'), exist_ok=True)
    os.makedirs(os.path.join(valDir, 'labels'), exist_ok=True)

    # 获取数据集中所有文件的列表
    file_list = os.listdir(srcDir)
    random.shuffle(file_list)

    # 根据划分比例计算训练集和验证集的边界索引
    split_index = int(len(file_list) * split_ratio)
    train_files = file_list[:split_index]
    val_files = file_list[split_index:]

    # 将训练集数据移动到相应文件夹
    for file in train_files:
        if file.endswith('.jpg'):
            img_src = os.path.join(srcDir, file)
            label_src = os.path.join(srcDir, file[:-4] + '.txt')
            shutil.copy(img_src, os.path.join(trainDir, 'images', file))
            shutil.copy(label_src, os.path.join(trainDir, 'labels', file[:-4] + '.txt'))

    # 将验证集数据移动到相应文件夹
    for file in val_files:
        if file.endswith('.jpg'):
            img_src = os.path.join(srcDir, file)
            label_src = os.path.join(srcDir, file[:-4] + '.txt')
            shutil.copy(img_src, os.path.join(valDir, 'images', file))
            shutil.copy(label_src, os.path.join(valDir, 'labels', file[:-4] + '.txt'))

if __name__ == '__main__':
    # 输入文件夹路径
    srcDir = r"E:\LR\handian20240801\imgs"
    trainDir = r'E:\LR\handian20240801\train'
    valDir = r'E:\LR\handian20240801\val'

    # 调用函数划分数据集
    split_dataset(srcDir, trainDir, valDir)
